from concurrent.futures import ThreadPoolExecutor
from xml.dom.minidom import parse
from xml.dom import minidom
import shutil
import psycopg2
import datetime
import zipfile
import paramiko
from scp import SCPClient
from osgeo import gdal
from osgeo import gdalconst
import os
import re

# get sentinel2 tile list
def get_tiles(
    province=None, 
    cityname=None, 
    countyname=None, 
    tiles=[]):
    con = psycopg2.connect(
    database="sentinel2",
    user="postgres",
    password="12345",
    host="192.168.2.172",
    port="5432"
    )
    cur = con.cursor()
    if province is not None:
        sql = f"""
        SELECT "中国有效哨兵格网分布wgs84".mgrs
        FROM china_sy
        JOIN "中国有效哨兵格网分布wgs84" 
        on st_intersects(china_sy.geom,"中国有效哨兵格网分布wgs84".geom)
        WHERE  china_sy.qhmc = '{province}' 
        """
    elif cityname is not None:
        sql = f"""
        SELECT "中国有效哨兵格网分布wgs84".mgrs
        FROM china_dsy
        JOIN "中国有效哨兵格网分布wgs84" 
        on st_intersects(china_dsy.geom,"中国有效哨兵格网分布wgs84".geom)
        WHERE  china_dsy.qhmc = '{cityname}' 
        """
    elif countyname is not None:
        sql = f"""
        SELECT "中国有效哨兵格网分布wgs84".mgrs
        FROM china_xy
        JOIN "中国有效哨兵格网分布wgs84" 
        on st_intersects(china_xy.geom,"中国有效哨兵格网分布wgs84".geom)
        WHERE  china_xy.qhmc = '{countyname}' 
        """
    cur.execute(sql)
    con.commit()
    tiles.extend([record[0] for record in cur.fetchall()])
    tiles = [tile.upper() for tile in tiles]
    cur.close()
    con.close()
    return tiles

# raster to array
def get_raster_band_array(raster_path, shape):
	""""""
	dataset = gdal.Open(raster_path, 1)
	(buf_xsize, buf_ysize) = shape
	xsize = dataset.RasterXSize
	ysize = dataset.RasterYSize
	array = dataset.ReadAsArray(
		xoff=0,
		yoff=0,
		xsize=xsize,
		ysize=ysize,
		buf_xsize=buf_xsize,
		buf_ysize=buf_ysize,
		resample_alg=gdalconst.GRIORA_NearestNeighbour,
		interleave='band')
	return array

# data to raster file
def get_raster_file(raster_path, arrays, attribute):
	""""""
	driver = gdal.GetDriverByName("GTiff")
	dataset = driver.Create(
		raster_path,
		xsize=attribute["RasterXSize"],
		ysize=attribute["RasterYSize"],
		bands=attribute["RasterCount"],
		eType=attribute["DataType"]
	)
	dataset.SetSpatialRef(attribute["SpatialRef"])
	dataset.SetGeoTransform(attribute["GeoTransform"])
	nodata = 0 if attribute["NoData"] is None else attribute["NoData"]
	for index in range(1, attribute["RasterCount"] + 1):
		band = dataset.GetRasterBand(index)
		band.WriteArray(arrays[index-1])
		band.SetNoDataValue(nodata)
	return

# raster to attribute
def get_raster_attribute(raster_path):
	""""""
	dataset = gdal.Open(raster_path)
	attributes = {
		"RasterXSize": dataset.__getattr__("RasterXSize"),
		"RasterYSize": dataset.__getattr__("RasterYSize"),
		"RasterCount": dataset.__getattr__("RasterCount"),
		"SpatialRef": dataset.GetSpatialRef(),
		"GeoTransform": dataset.GetGeoTransform(),
		"DataType": dataset.GetRasterBand(1).DataType,
		"NoData": dataset.GetRasterBand(1).GetNoDataValue(),
	}
	return attributes

# read clouds of xml
def get_l1c_cloud(zip_path):
	""""""
	select_node = [
		{
			"tag_name": "n1:Quality_Indicators_Info",
			"attribute": None
		},
		{
			"tag_name": "Cloud_Coverage_Assessment",
			"attribute": None
		}
	]
	xml_namere = "^S2.*\.SAFE/MTD_MSIL1C\.xml$"
	zip_file = zipfile.ZipFile(zip_path, "r")
	for file_info in zip_file.infolist():
		if re.match(xml_namere, file_info.filename):
			xml_path = zip_file.open(file_info.filename, "r")
			datas = read_xml_data(xml_path, select_node)
			cloud = round(float(datas[0]))
			cloud = f"0{cloud}" if len(f"{cloud}") == 1 else f"{cloud}"
	return cloud

# read xml
def read_xml_data(file_path, select_node):
	""""""
	domTree = minidom.parse(file_path)
	collection = domTree.documentElement
	select_element = [collection]
	for node_dict in select_node:
		tag_name = node_dict["tag_name"]
		if node_dict["attribute"] is None:
			select_element = [
				element_sub
				for element in select_element
			    for element_sub in element.getElementsByTagName(tag_name)
			]
		else:
			attribute_name = node_dict["attribute"]["name"]
			attribute_value = node_dict["attribute"]["value"]
			select_element = [
				element_sub
				for element in select_element
				for element_sub in element.getElementsByTagName(tag_name)
				if element_sub.getAttribute(attribute_name) == attribute_value
			]
	datas = [
		element.firstChild.data
		for element in select_element
	]
	return datas

# single band rasters to mutil band raster
def get_mutil_band_raster(source_paths, target_path):
	""""""
	count = len(source_paths)
	max_workers = count

	# 多线程的读取栅格文件属性
	with ThreadPoolExecutor(max_workers=max_workers) as executor:
		attributes = executor.map(get_raster_attribute, source_paths)

	# 获得模板属性
	attributes = list(attributes)
	pixelsizes = [attribute["GeoTransform"][1] for attribute in attributes]
	attribute = attributes[pixelsizes.index(min(pixelsizes))]
	shape = (attribute["RasterXSize"], attribute["RasterYSize"])
	attribute.update({"RasterCount": count})

	# 多线程的读取栅格文件数组
	# shapes = [shape] * count
	# with ThreadPoolExecutor(max_workers=5) as executor:
	# 	arrays = executor.map(get_raster_band_array, source_paths, shapes)

	arrays = [
		get_raster_band_array(source_path, shape)
		for source_path in source_paths
	]

	# 数据写入栅格
	arrays = list(arrays)
	get_raster_file(target_path, arrays, attribute)
	return

# get l1c raster save folder
def get_folder( 
    year, month, province=None, city=None, county=None):
    con = psycopg2.connect(
    database="sentinel2",
    user="postgres",
    password="12345",
    host="192.168.2.172",
    port="5432"
    )
    cur = con.cursor()
    if city is not None:
        sql = f"""
        SELECT qhmc
        FROM china_sy
        WHERE qhdm in (
        SELECT sjxzqdm
        FROM china_dsy
        WHERE qhmc='{city}'
        )
        """
        cur.execute(sql)
        con.commit()
        province = cur.fetchall()[0][0]
    elif county is not None:
        sql = f"""
        SELECT qhmc
        FROM china_sy
        WHERE qhdm IN (
        SELECT sjxzqdm
        FROM china_dsy
        WHERE qhdm IN (
        SELECT sjxzqdm
        FROM china_xy
        WHERE qhmc = '{county}'))
        """
        cur.execute(sql)
        con.commit()
        province = cur.fetchall()[0][0]
    else:
        pass
    cur.close()
    con.close()
    folder = os.path.join(
            "/home/遥感技术部/遥感数据/A01影像", 
            f"{province}",
            f"{year}年",
            f"{province}{year}hc",
            f"{month}月"
        )
    if os.path.exists(folder):
        print(folder)
    return folder

# select sentinel2 l1c zip file path
def select_zips(tile, date, cloudmax=100):
    con = psycopg2.connect(
    database="sentinel2",
    user="postgres",
    password="12345",
    host="192.168.2.172",
    port="5432"
    )
    cur = con.cursor()
    sql = f"""
    SELECT path
    FROM l1c_zip_info 
    WHERE to_char(shot_date, 'yyyyMMdd') like '{date}%'
    AND tile = '{tile.upper()}' AND clouds <= {cloudmax}
    """
    cur.execute(sql)
    con.commit()
    paths = [record[0] for record in cur.fetchall()]
    cur.close()
    con.close()
    return paths

# unzip zip to file
def unzip_file(zip_path, nameres, save_dir):
	""""""
	paths = []
	zipFile = zipfile.ZipFile(zip_path, "r")
	for namere in nameres:
		for info in zipFile.infolist():
			if re.match(namere, os.path.basename(info.filename)):
				path = zipFile.extract(info.filename, save_dir)
				paths.append(path)
	zipFile.close()
	return paths

# scp remote file
def scp_file(file_path, folder):
    """"""
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy)
    ssh.connect(
        hostname="192.168.2.172",
        port=22, 
        username="root", 
        password="DELL123dell")
    #  通过scp上传或下载服务器文件
    scpclient = SCPClient(ssh.get_transport(), socket_timeout=15.0)
    file_name = os.path.basename(file_path)
    dest_path = os.path.join(folder, file_name)
    if not os.path.exists(file_path):
        scpclient.get(file_path, dest_path)
    else:
        scpclient.put(file_path, dest_path)
    return dest_path

# get sentinel2 l1c raster file
def get_raster(zip_path, save_dir, proj_dir=None, bands=None):
    if proj_dir is not None:
        os.environ['PROJ_LIB'] = proj_dir
    if bands is None:
        bands = [2, 3, 4, 5, 6, 7, 8, "8A", 12]
    cloud = get_l1c_cloud(zip_path)
    zip_name = os.path.basename(zip_path)
    info = zip_name.split("_")
    tile = info[5][1:]
    date = info[2][:8]
    time = info[6][9:15]
    raster_name = f"L1C_{tile}_{date}-{time}_{cloud}.tif"
    raster_path = os.path.join(save_dir, raster_name)
    if os.path.exists(raster_path):
        return raster_path
    nameres = [
		f"T.*B{'0'*(len(f'{index}') < 2)}{index}\.jp2"
		for index in bands
	]
    band_paths = unzip_file(zip_path, nameres, save_dir)
    cloud = get_l1c_cloud(zip_path)
    os.remove(zip_path)
    get_mutil_band_raster(band_paths, raster_path)
    shutil.rmtree(os.path.join(zip_path.replace(".zip",".SAFE")))
    return raster_path

# start
def main():
    
	return

if __name__ == "__main__":
    
    tiles =get_tiles(province="山东省")
    date = "202203"
    cloud = 100
    for tile in tiles:
        zip_paths = select_zips(tile, date, cloud)
        for zip_path in zip_paths:
            print(zip_path)